{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\nTrain a neural network potential\n================================\n\nIn this tutorial, we train a neural network (NN) potential for silicon\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We are going to fit the NN potential to a training set of energies and forces from\ncompressed and stretched diamond silicon structures (the same training set used in\n`tut_kim_sw`).\nDownload the training set :download:`Si_training_set.tar.gz <https://raw.githubusercontent.com/mjwen/kliff/master/examples/Si_training_set.tar.gz>`\nand extract the tarball: ``$ tar xzf Si_training_set.tar.gz``.\nThe data is stored in **extended xyz** format, and see `doc.dataset` for more\ninformation of this format.\n\n<div class=\"alert alert-danger\"><h4>Warning</h4><p>The ``Si_training_set`` is just a toy data set for the purpose to demonstrate how to\n    use KLIFF to train potentials. It should not be used to train any potential for real\n    simulations.</p></div>\n\nLet's first import the modules that will be used in this example.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from kliff.dataset import Dataset\nfrom kliff.descriptors import SymmetryFunction\nfrom kliff.models import NeuralNetwork\nfrom kliff.calculators import CalculatorTorch\nfrom kliff import nn\nfrom kliff.loss import Loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Model\n-----\n\nFor a NN model, we need to specify the descriptor that transforms atomic environment\ninformation to the fingerprints, which the NN model uses as the input. Here, we use the\nsymmetry functions proposed by Behler and coworkers.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "descriptor = SymmetryFunction(\n    cut_name=\"cos\", cut_dists={\"Si-Si\": 5.0}, hyperparams=\"set30\", normalize=True\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The ``cut_name`` and ``cut_dists`` tell the descriptor what type of cutoff function to\nuse and what the cutoff distances are. ``hyperparams`` specifies the set of\nhyperparameters used in the symmetry function descriptor. If you prefer, you can provide\na dictionary of your own hyperparameters. And finally, ``normalize`` informs that the\ngenerated fingerprints should be normalized by first subtracting the mean and then\ndividing the standard deviation. This normalization typically makes it easier to\noptimize NN model.\n\nWe can then build the NN model on top of the descriptor.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "N1 = 10\nN2 = 10\nmodel = NeuralNetwork(descriptor)\nmodel.add_layers(\n    # first hidden layer\n    nn.Linear(descriptor.get_size(), N1),\n    nn.Tanh(),\n    # second hidden layer\n    nn.Linear(N1, N2),\n    nn.Tanh(),\n    # output layer\n    nn.Linear(N2, 1),\n)\nmodel.set_save_metadata(prefix=\"./kliff_saved_model\", start=5, frequency=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In the above code, we build a NN model with an input layer, two hidden layer, and an\noutput layer. The ``descriptor`` carries the information of the input layer, so it is\nnot needed to be specified explicitly. For each hidden layer, we first do a linear\ntransformation using ``nn.Linear(size_in, size_out)`` (essentially carrying out $y\n= xW+b$, where $W$ is the weight matrix of size ``size_in`` by ``size_out``, and\n$b$ is a vector of size ``size_out``. Then we apply the hyperbolic tangent\nactivation function ``nn.Tanh()`` to the output of the Linear layer (i.e. $y$) so\nas to add the nonlinearity. We use a Linear layer for the output layer as well, but\nunlike the hidden layer, no activation function is applied here. The input size\n``size_in`` of the first hidden layer must be the size of the descriptor, which is\nobtained using ``descriptor.get_size()``. For all other layers (hidden or output), the\ninput size must be equal to the output size of the previous layer. The ``out_size`` of\nthe output layer must be 1 such that the output of the NN model gives the energy of the\natom.\n\nThe ``set_save_metadata`` function call informs where to save intermediate models during\nthe optimization (discussed below), and what the starting epoch and how often to save\nthe model.\n\n\nTraining set and calculator\n---------------------------\n\nThe training set and the calculator are the same as explained in `tut_kim_sw`. The\nonly difference is that we need to use the\n:mod:`~kliff.calculators.CalculatorTorch()`, which is targeted for the NN model.\nAlso, its ``create()`` method takes an argument ``reuse`` to inform whether to reuse the\nfingerprints generated from the descriptor if it is present.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# training set\ndataset_name = \"Si_training_set/varying_alat\"\ntset = Dataset()\ntset.read(dataset_name)\nconfigs = tset.get_configs()\n\n# calculator\ncalc = CalculatorTorch(model)\ncalc.create(configs, reuse=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Loss function\n-------------\n\nKLIFF uses a loss function to quantify the difference between the training data and\npotential predictions and uses minimization algorithms to reduce the loss as much as\npossible. In the following code snippet, we create a loss function that uses the\n``Adam`` optimizer to minimize it. The Adam optimizer supports minimization using\n`mini-batches` of data, and here we use ``100`` configurations in each minimization step\n(the training set has a total of 400 configurations as can be seen above), and run\nthrough the training set for ``10`` epochs. The learning rate ``lr`` used here is\n``0.001``, and typically, one may need to play with this to find an acceptable one that\ndrives the loss down in a reasonable time.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "loss = Loss(calc, residual_data={\"forces_weight\": 0.3})\nresult = loss.minimize(method=\"Adam\", num_epochs=10, batch_size=100, lr=0.001)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can save the trained model to disk, and later can load it back if we want. We can\nalso write the trained model to a KIM model such that it can be used in other simulation\ncodes such as LAMMPS via the KIM API.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model.save(\"./final_model.pkl\")\nloss.save_optimizer_stat(\"./optimizer_stat.pkl\")\n\nmodel.write_kim_model()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>Now we have trained an NN for a single specie Si. If you have multiple species in\n   your system and want to use different parameters for different species,\n   take a look at the SiC_ example.</p></div>\n\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}